#include "DCBEstimator.h"
#include <iostream>
#include <set>
#include <cmath>
#include <algorithm>

DCBEstimator::DCBEstimator() {
    // 初始化
}

DCBEstimator::~DCBEstimator() {
    // 清理资源
}

void DCBEstimator::setObservations(const std::vector<VTECObservation>& vtecObs) {
    vtecObservations_ = vtecObs;
}

void DCBEstimator::setSphericalHarmonicModel(const SphericalHarmonic& shModel) {
    shModel_ = shModel;
}

void DCBEstimator::estimateDCB() {
    // 1. 提取所有卫星和接收机ID
    std::set<std::string> satellites;
    std::set<std::string> receivers;
    
    for (const auto& obs : vtecObservations_) {
        satellites.insert(obs.satellite);
        receivers.insert(obs.station);
    }
    
    int numSatellites = satellites.size();
    int numReceivers = receivers.size();
    
    std::cout << "DCB估计: " << numSatellites << " 颗卫星, " << numReceivers << " 个接收机" << std::endl;
    
    // 2. 构建卫星和接收机ID到索引的映射
    std::map<std::string, int> satIndices;
    std::map<std::string, int> recIndices;
    
    int idx = 0;
    for (const auto& sat : satellites) {
        satIndices[sat] = idx++;
    }
    
    idx = 0;
    for (const auto& rec : receivers) {
        recIndices[rec] = idx++;
    }
    
    // 3. 构建设计矩阵和观测向量
    torch::Tensor designMatrix = buildDCBDesignMatrix();
    
    // 准备观测向量: STEC - VTEC/MF
    auto options = torch::TensorOptions().dtype(torch::kDouble);
    torch::Tensor obsVector = torch::zeros({static_cast<long int>(vtecObservations_.size())}, options);
    
    for (size_t i = 0; i < vtecObservations_.size(); ++i) {
        const auto& obs = vtecObservations_[i];
        double predictedVTEC = shModel_.predictVTEC(obs.latitude, obs.longitude, obs.time);
        double stec = obs.stec;
        double mf = obs.mappingFunction;
        
        // 观测方程: STEC - VTEC/MF = DCB_sat + DCB_rec
        obsVector[static_cast<long int>(i)] = stec - predictedVTEC / mf;
    }
    
    // 4. 添加约束条件: 使卫星DCB的和为0
    applyZeroMeanConstraint(designMatrix, obsVector);
    
    // 5. 修改求解部分，使用SVD分解
    try {
        // Fix the torch::linalg::lstsq call by adding the required driver (or nullptr) and 'gels' arguments
        // lstsq(input, other, driver, method)
        auto lstsq_result = torch::linalg::lstsq(
            designMatrix, 
            obsVector.unsqueeze(1),
            std::nullopt,  // Use default rcond value
            "gels"  // Use default method
        );
        
        // 从返回的元组中提取solution
        torch::Tensor solution = std::get<0>(lstsq_result);
        
        // 检查解的合理性
        double maxDCB = solution.abs().max().item().toDouble();
        if (maxDCB > 50.0) {  // DCB通常不会超过50ns
            throw std::runtime_error("DCB值异常大，切换到SVD求解");
        }
        
        // 提取DCB值
        extractDCBValues(solution);
    }
    catch (const std::exception& e) {
        std::cout << "使用SVD求解DCB: " << e.what() << std::endl;
        
        // Fix the torch::linalg::svd call by adding the required parameter
        // svd(input, full_matrices, driver)
        auto [U, S, Vh] = torch::linalg::svd(
            designMatrix, 
            true,  // full_matrices
            std::nullopt  // Use default driver
        );
        
        // Fix the .item<double>() syntax
        double max_s = torch::max(S).item().toDouble();
        double threshold = max_s * 1e-10;  // 设置阈值为最大奇异值的1e-10倍
        
        torch::Tensor Sinv = torch::zeros_like(S);
        for (int i = 0; i < S.size(0); i++) {
            if (S[i].item().toDouble() > threshold) {
                Sinv[i] = 1.0 / S[i].item().toDouble();
            }
        }
        
        // 计算解：V * S⁻¹ * U.T * b
        auto V = Vh.transpose(0, 1);
        auto USinv = U * Sinv.unsqueeze(1);
        torch::Tensor solution = torch::matmul(V, torch::matmul(USinv.transpose(0, 1), obsVector.unsqueeze(1)));
        
        // 提取DCB值
        satellites_.clear();
        receivers_.clear();
        
        for (const auto& sat : satellites) {
            satellites_.insert(sat);
        }
        
        for (const auto& rec : receivers) {
            receivers_.insert(rec);
        }
        
        // 提取DCB值
        satelliteDCBs_.clear();
        receiverDCBs_.clear();
        
        for (const auto& [satId, index] : satIndices) {
            DCBResult result;
            result.id = satId;
            result.dcb = solution[index].item().toDouble();
            
            // 估计标准差
            double sigma = 0.0;
            for (int i = 0; i < S.size(0); i++) {
                if (S[i].item().toDouble() > threshold) {
                    double v_element = V[index][i].item().toDouble();
                    sigma += (v_element * v_element) / (S[i].item().toDouble() * S[i].item().toDouble());
                }
            }
            sigma = std::sqrt(sigma);
            
            result.dcbSigma = sigma;
            satelliteDCBs_.push_back(result);
        }
        
        for (const auto& [recId, index] : recIndices) {
            DCBResult result;
            result.id = recId;
            result.dcb = solution[numSatellites + index].item().toDouble();
            
            // 估计标准差 (类似于卫星DCB)
            double sigma = 0.0;
            for (int i = 0; i < S.size(0); i++) {
                if (S[i].item().toDouble() > threshold) {
                    double v_element = V[numSatellites + index][i].item().toDouble();
                    sigma += (v_element * v_element) / (S[i].item().toDouble() * S[i].item().toDouble());
                }
            }
            sigma = std::sqrt(sigma);
            
            result.dcbSigma = sigma;
            receiverDCBs_.push_back(result);
        }
    }
    
    std::cout << "DCB估计完成" << std::endl;
}

// 联合估计VTEC球谐系数和DCB
void DCBEstimator::estimateVTECAndDCB(int maxDegree) {
    // 1. 提取所有卫星和接收机ID
    std::set<std::string> satellites;
    std::set<std::string> receivers;
    
    for (const auto& obs : vtecObservations_) {
        satellites.insert(obs.satellite);
        receivers.insert(obs.station);
    }
    
    satellites_ = satellites;
    receivers_ = receivers;
    
    int numSatellites = satellites.size();
    int numReceivers = receivers.size();
    
    std::cout << "联合估计VTEC与DCB: " << numSatellites << " 颗卫星, " << numReceivers << " 个接收机" << std::endl;
    
    // 2. 构建联合估计的设计矩阵
    torch::Tensor designMatrix = buildJointDesignMatrix(maxDegree);
    
    // 计算球谐模型的系数数量
    int numSHCoeffs = (maxDegree + 1) * (maxDegree + 1);
    
    // 3. 准备观测向量: STEC
    auto options = torch::TensorOptions().dtype(torch::kDouble);
    torch::Tensor obsVector = torch::zeros({static_cast<long int>(vtecObservations_.size())}, options);
    torch::Tensor weightVector = torch::ones({static_cast<long int>(vtecObservations_.size())}, options);
    
    for (size_t i = 0; i < vtecObservations_.size(); ++i) {
        const auto& obs = vtecObservations_[i];
        obsVector[static_cast<long int>(i)] = obs.stec;
        
        // 使用高度角的正弦值作为权重 (高度角越高，权重越大)
        double elevation = obs.elevation * M_PI / 180.0; // 转换为弧度
        weightVector[static_cast<long int>(i)] = std::sin(elevation);
    }
    
    // 4. 添加约束条件: 使接收机DCB的和为0
    applyZeroMeanConstraint(designMatrix, obsVector);
    
    // 5. 应用加权最小二乘
    // 计算加权设计矩阵和观测向量
    torch::Tensor sqrtWeights = torch::sqrt(weightVector);
    
    torch::Tensor weightedDesignMatrix = designMatrix.clone();
    torch::Tensor weightedObsVector = obsVector.clone();
    
    for (int i = 0; i < designMatrix.size(0); i++) {
        weightedDesignMatrix[i] = designMatrix[i] * sqrtWeights[i];
        weightedObsVector[i] = obsVector[i] * sqrtWeights[i];
    }
    
    // 6. 求解联合方程
    try {
        auto lstsq_result = torch::linalg::lstsq(
            weightedDesignMatrix, 
            weightedObsVector.unsqueeze(1),
            std::nullopt,  // Use default rcond value
            "gels"  // Use default method
        );
        
        torch::Tensor solution = std::get<0>(lstsq_result);
        
        // 从解中提取球谐系数和DCB值
        extractDCBAndSHValues(solution, numSHCoeffs);
    }
    catch (const std::exception& e) {
        std::cout << "使用SVD求解联合方程: " << e.what() << std::endl;
        
        // 使用SVD求解
        auto [U, S, Vh] = torch::linalg::svd(
            weightedDesignMatrix, 
            true,  // full_matrices
            std::nullopt  // Use default driver
        );
        
        double max_s = torch::max(S).item().toDouble();
        double threshold = max_s * 1e-10;  // 设置阈值为最大奇异值的1e-10倍
        
        torch::Tensor Sinv = torch::zeros_like(S);
        for (int i = 0; i < S.size(0); i++) {
            if (S[i].item().toDouble() > threshold) {
                Sinv[i] = 1.0 / S[i].item().toDouble();
            }
        }
        
        // 计算解：V * S⁻¹ * U.T * b
        auto V = Vh.transpose(0, 1);
        auto USinv = U * Sinv.unsqueeze(1);
        torch::Tensor solution = torch::matmul(V, torch::matmul(USinv.transpose(0, 1), weightedObsVector.unsqueeze(1)));
        
        // 从解中提取球谐系数和DCB值
        extractDCBAndSHValues(solution, numSHCoeffs);
    }
    
    std::cout << "联合估计VTEC与DCB完成" << std::endl;
}

// 构建联合估计的设计矩阵
torch::Tensor DCBEstimator::buildJointDesignMatrix(int maxDegree) const {
    // 1. 构建卫星和接收机ID到索引的映射
    std::map<std::string, int> satIndices;
    std::map<std::string, int> recIndices;
    
    int idx = 0;
    for (const auto& sat : satellites_) {
        satIndices[sat] = idx++;
    }
    
    idx = 0;
    for (const auto& rec : receivers_) {
        recIndices[rec] = idx++;
    }
    
    int numSatellites = satellites_.size();
    int numReceivers = receivers_.size();
    int numObservations = vtecObservations_.size();
    
    // 2. 计算球谐系数的数量
    int numSHCoeffs = (maxDegree + 1) * (maxDegree + 1);
    
    // 3. 构建联合设计矩阵
    auto options = torch::TensorOptions().dtype(torch::kDouble);
    int numColumns = numSHCoeffs + numSatellites + numReceivers;
    torch::Tensor designMatrix = torch::zeros({numObservations, numColumns}, options);
    
    for (int i = 0; i < numObservations; ++i) {
        const auto& obs = vtecObservations_[i];
        double mf = obs.mappingFunction;
        
        // 填充球谐函数部分
        // 转换为余纬度和弧度
        double colatitude = (90.0 - obs.latitude) * M_PI / 180.0;
        double longitude = obs.longitude * M_PI / 180.0;
        
        int col = 0;
        for (int n = 0; n <= maxDegree; ++n) {
            for (int m = 0; m <= n; ++m) {
                // 计算勒让德函数
                double pnm = Utils::normalizedLegendre(n, m, std::cos(colatitude));
                
                // 余弦项
                designMatrix[i][col] = (pnm * std::cos(m * longitude)) / mf;
                col++;
                
                // 正弦项 (m > 0时)
                if (m > 0) {
                    designMatrix[i][col] = (pnm * std::sin(m * longitude)) / mf;
                    col++;
                }
            }
        }
        
        // 填充DCB部分
        int satIdx = satIndices[obs.satellite];
        designMatrix[i][numSHCoeffs + satIdx] = 1.0;
        
        int recIdx = recIndices[obs.station];
        designMatrix[i][numSHCoeffs + numSatellites + recIdx] = 1.0;
    }
    
    return designMatrix;
}

// 从解向量中提取DCB值和球谐系数
void DCBEstimator::extractDCBAndSHValues(const torch::Tensor& solution, int numSHCoeffs) {
    // 提取球谐系数
    shCoefficients_.clear();
    for (int i = 0; i < numSHCoeffs; i++) {
        shCoefficients_.push_back(solution[i].item().toDouble());
    }
    
    // 构建ID到索引的映射
    std::map<std::string, int> satIndices;
    std::map<std::string, int> recIndices;
    
    int idx = 0;
    for (const auto& sat : satellites_) {
        satIndices[sat] = idx++;
    }
    
    idx = 0;
    for (const auto& rec : receivers_) {
        recIndices[rec] = idx++;
    }
    
    int numSatellites = satellites_.size();
    
    // 提取卫星DCB
    satelliteDCBs_.clear();
    for (const auto& [satId, index] : satIndices) {
        DCBResult result;
        result.id = satId;
        result.dcb = solution[numSHCoeffs + index].item().toDouble();
        result.dcbSigma = 0.1; // 设置一个默认标准差值
        satelliteDCBs_.push_back(result);
    }
    
    // 提取接收机DCB
    receiverDCBs_.clear();
    for (const auto& [recId, index] : recIndices) {
        DCBResult result;
        result.id = recId;
        result.dcb = solution[numSHCoeffs + numSatellites + index].item().toDouble();
        result.dcbSigma = 0.1; // 设置一个默认标准差值
        receiverDCBs_.push_back(result);
    }
}

std::vector<DCBResult> DCBEstimator::getSatelliteDCBs() const {
    return satelliteDCBs_;
}

std::vector<DCBResult> DCBEstimator::getReceiverDCBs() const {
    return receiverDCBs_;
}

torch::Tensor DCBEstimator::buildDCBDesignMatrix() const {
    // 提取所有卫星和接收机ID
    std::set<std::string> satellites;
    std::set<std::string> receivers;
    
    for (const auto& obs : vtecObservations_) {
        satellites.insert(obs.satellite);
        receivers.insert(obs.station);
    }
    
    int numSatellites = satellites.size();
    int numReceivers = receivers.size();
    int numParameters = numSatellites + numReceivers;
    int numObservations = vtecObservations_.size();
    
    // 构建卫星和接收机ID到索引的映射
    std::map<std::string, int> satIndices;
    std::map<std::string, int> recIndices;
    
    int idx = 0;
    for (const auto& sat : satellites) {
        satIndices[sat] = idx++;
    }
    
    idx = 0;
    for (const auto& rec : receivers) {
        recIndices[rec] = idx++;
    }
    
    // 构建设计矩阵
    auto options = torch::TensorOptions().dtype(torch::kDouble);
    torch::Tensor designMatrix = torch::zeros({numObservations, numParameters}, options);
    
    for (int i = 0; i < numObservations; ++i) {
        const auto& obs = vtecObservations_[i];
        
        // 设置卫星DCB列
        int satIdx = satIndices[obs.satellite];
        designMatrix[i][satIdx] = 1.0;
        
        // 设置接收机DCB列
        int recIdx = recIndices[obs.station];
        designMatrix[i][numSatellites + recIdx] = 1.0;
    }
    
    return designMatrix;
}

void DCBEstimator::applyZeroMeanConstraint(torch::Tensor& designMatrix, torch::Tensor& observationVector) {
    // 获取矩阵尺寸
    int numParams = designMatrix.size(1);
    
    // 提取所有卫星和接收机ID
    std::set<std::string> satellites;
    std::set<std::string> receivers;
    
    for (const auto& obs : vtecObservations_) {
        satellites.insert(obs.satellite);
        receivers.insert(obs.station);
    }
    int numSatellites = satellites.size();
    int numReceivers = receivers.size();
    
    // 添加约束方程: 修改为所有接收机DCB的和为0，而不是卫星DCB
    auto options = torch::TensorOptions().dtype(torch::kDouble);
    torch::Tensor constraintRow = torch::zeros({1, numParams}, options);
    
    // 设置卫星DCB项为0，接收机DCB项为1
    for (int i = 0; i < numReceivers; ++i) {
        constraintRow[0][numSatellites + i] = 1.0;
    }
    
    // 扩展设计矩阵
    torch::Tensor newDesignMatrix = torch::cat({designMatrix, constraintRow}, 0);
    
    // 扩展观测向量，添加约束值0
    torch::Tensor constraintValue = torch::zeros({1}, options);
    torch::Tensor newObsVector = torch::cat({observationVector, constraintValue}, 0);
    
    // 更新原始引用
    designMatrix = newDesignMatrix;
    observationVector = newObsVector;
}

void DCBEstimator::extractDCBValues(const torch::Tensor& solution) {
    // 提取所有卫星和接收机ID
    std::set<std::string> satellites;
    std::set<std::string> receivers;
    
    for (const auto& obs : vtecObservations_) {
        satellites.insert(obs.satellite);
        receivers.insert(obs.station);
    }
    
    satellites_ = satellites;
    receivers_ = receivers;
    
    // 提取DCB解算结果并存储
    int numSatellites = satellites.size();
    int numReceivers = receivers.size();
    
    // 确保解的大小符合预期
    if (solution.size(0) != numSatellites + numReceivers) {
        throw std::runtime_error("DCB解向量大小与卫星和接收机数量不匹配");
    }
    
    // 清空之前的结果
    satelliteDCBs_.clear();
    receiverDCBs_.clear();
    
    // 创建卫星和接收机的索引映射
    std::map<std::string, int> satIndices;
    std::map<std::string, int> recIndices;
    
    int idx = 0;
    for (const auto& sat : satellites) {
        satIndices[sat] = idx++;
    }
    
    idx = 0;
    for (const auto& rec : receivers) {
        recIndices[rec] = idx++;
    }
    
    // 提取卫星DCB
    for (const auto& [satId, index] : satIndices) {
        DCBResult result;
        result.id = satId;
        result.dcb = solution[index].item().toDouble();
        result.dcbSigma = 0.1; // 设置一个默认值，可以根据实际情况计算
        satelliteDCBs_.push_back(result);
    }
    
    // 提取接收机DCB
    for (const auto& [recId, index] : recIndices) {
        DCBResult result;
        result.id = recId;
        result.dcb = solution[numSatellites + index].item().toDouble();
        result.dcbSigma = 0.1; // 设置一个默认值，可以根据实际情况计算
        receiverDCBs_.push_back(result);
    }
}

std::map<std::string, double> DCBEstimator::getSatelliteDCBMap() const {
    std::map<std::string, double> dcbMap;
    
    for (const auto& dcb : satelliteDCBs_) {
        dcbMap[dcb.id] = dcb.dcb;
    }
    
    return dcbMap;
}

std::map<std::string, double> DCBEstimator::getReceiverDCBMap() const {
    std::map<std::string, double> dcbMap;
    
    for (const auto& dcb : receiverDCBs_) {
        dcbMap[dcb.id] = dcb.dcb;
    }
    
    return dcbMap;
}

std::vector<double> DCBEstimator::getSphericalHarmonicCoefficients() const {
    return shCoefficients_;
} 